import torch
import torch.nn as nn
import torch.nn.functional as F


class ResidualBlcok(nn.Module):
    """
    残差块:
    """
    def __init__(self,input_channels, output_channels, use_conv=False, stride=1):
        """
        残差块
        :param input_channels: 输入通道数
        :param output_channels: 输出通道数
        :param use_conv: 是否使用卷积层对输入通道数进行转换
        :param stride: 步长
        """
        super().__init__()
        self.covn1 = nn.Conv2d(input_channels, output_channels, kernel_size=3, stride=stride, padding=1)
        self.bn1 = nn.BatchNorm2d(output_channels)
        self.relu = nn.ReLU(inplace=True)
        self.covn2 = nn.Conv2d(output_channels, output_channels, kernel_size=3, stride=1, padding=1)
        self.bn2 = nn.BatchNorm2d(output_channels)
        if use_conv:
            self.conv3 = nn.Conv2d(input_channels, output_channels, kernel_size=1, stride=stride)
        else:
            self.conv3 = None
    
    def forward(self, X):
        """
        前向传播
        :param X: 输入
        :return: 输出
        """
        Y = self.relu(self.bn1(self.covn1(X)))
        Y = self.bn2(self.covn2(Y))
        if self.conv3:
            X = self.conv3(X)
        Y += X
        return F.relu(Y)


class Resnet18(nn.Module):
    """
    Resnet18模型
    """
    def __init__(self, num_classes=100):
        super().__init__()
        self.b1 = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        )
        self.b2 = self.resnet_block(64, 64, 2, first_block=True)
        self.b3 = self.resnet_block(64, 128, 2)
        self.b4 = self.resnet_block(128, 256, 2)
        self.b5 = self.resnet_block(256, 512, 2)
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.flatten = nn.Flatten()
        self.fc = nn.Linear(512, num_classes)
    
    def resnet_block(self, input_channels, output_channels, num_residuals, first_block=False):
        """
        残差块
        :param input_channels: 输入通道数
        :param output_channels: 输出通道数
        :param num_residuals: 残差块数量
        :param first_block: 是否为第一个残差块
        :return: 残差块
        """
        blk = []
        for i in range(num_residuals):
            if i == 0 and not first_block:
                blk.append(ResidualBlcok(input_channels, output_channels, use_conv=True, stride=2))
            else:
                blk.append(ResidualBlcok(output_channels, output_channels))
        return nn.Sequential(*blk)
    
    def forward(self, X):
        Y = self.b1(X)
        Y = self.b2(Y)
        Y = self.b3(Y)
        Y = self.b4(Y)
        Y = self.b5(Y)
        Y = self.avgpool(Y)
        Y = self.flatten(Y)
        Y = self.fc(Y)
        return Y


# X = torch.randn(1, 3, 224, 224)
# net = Resnet18()
# for layer in net.children():
#     X = layer(X)
#     print(layer.__class__.__name__, 'output shape:\t', X.shape)


# resnet 18 改 resnet 101
# 加载参数 name_parameters
# mobilenet


